package main

import (
	tr "tree"
)

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */

func pruneTree(root *tr.TreeNode) *tr.TreeNode {
	if root == nil {
		return nil
	}
	dfs(root)
	//对root的子树剪枝后，如果root是0叶子节点，则需要减掉本身
	if root.Val == 0 && root.Left == nil && root.Right == nil {
		return nil
	}
	return root
}

/*
root的所有节点有一个1 return true
全0 return false

并在后续遍历过程中剪掉全0的子树
*/
func dfs(root *tr.TreeNode) bool {
	if root == nil {
		return false
	}
	l := dfs(root.Left)
	r := dfs(root.Right)
	if !l {
		root.Left = nil
	}
	if !r {
		root.Right = nil
	}
	return root.Val == 1 || l || r
}

func main() {
	lst := tr.BuildList(1, nil, 0, 0, 1)
	root := tr.BuildTree(lst)
	res := pruneTree(root)
	tr.PrintList(tr.Tree2list(res))
}
